package rate_limiter

import (
	"context"
	"errors"
	"golang.org/x/time/rate"
	"time"
)

type TokenLimit struct {
	Limiter   *rate.Limiter
	IsSerial  bool
	WaitTimer time.Duration
}

func NewTokenLimit() *TokenLimit {
	return &TokenLimit{}
}

func (tl *TokenLimit) Init(limit, bucketSize int) {
	tl.Limiter = rate.NewLimiter(rate.Limit(limit), bucketSize)
}

func (tl *TokenLimit) CanNext(ctx context.Context) error {
	if t, ok := ctx.Deadline(); ok {
		reserve := tl.Limiter.Reserve()
		if reserve.OK() {
			return nil
		}
		if time.Now().Add(reserve.Delay()).Sub(t) > 0 {
			defer func() {
				reserve.Cancel() //归还令牌
			}()
			return errors.New("reserve time wait too large")
		}
	}

	if tl.IsSerial {
		ct, cancel := context.WithTimeout(ctx, tl.WaitTimer)
		defer cancel()
		if e := tl.Limiter.Wait(ct); e != nil {
			return e
		}
		return nil
	}

	if tl.Limiter.Allow() {
		return nil
	}

	return errors.New("limiter hidden it")
}
